# -*- coding: utf-8 -*-
# Author: Sun Jiawei
# E-mail: sunjiawei@tbea.com
#
"""
数据包含与电力系统建模相关的各种时间序列数据，即电价、电力消耗（负荷）以及风能和太阳能发电量和容量。
数据按国家、控制区或投标区汇总。地理覆盖范围包括欧盟和一些邻国。所有变量都以小时分辨率提供。
如果原始数据的分辨率较高（半小时或四分之一小时）,则会在单独的文件中提供.

脚本提供了对原始数据的可视化功能.
"""
import warnings
import pandas as pd
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # 用来正常显示中文标签


class DataSet:
    def __init__(self, csv_path: str) -> None:
        """Initialize a `DataSet` object."""
        self.data = pd.read_csv(csv_path)
        self.data['utc_timestamp'] = pd.to_datetime(self.data['utc_timestamp'])
        self.data = self.data.set_index('utc_timestamp')
        self.data['hour'] = self.data.index.hour
        self.data['minute'] = self.data.index.minute
        self.data['week'] = self.data.index.weekday + 1

    def get_date_range_data(self, start_year: str, start_month: str, start_day: str,
                            end_year: str, end_month: str, end_day: str) -> pd.DataFrame:
        """获取指定时间短内的数据.

        :param start_year: 开始年
        :param start_month: 开始月
        :param start_day: 开始日
        :param end_year: 终止年
        :param end_month: 终止月
        :param end_day: 终止日
        :return:
        """
        start = pd.Timestamp(start_year + '-' + start_month + '-' + start_day)
        end = pd.Timestamp(end_year + '-' + end_month + '-' + end_day)
        return self.data.truncate(before=start, after=end)

    def get_panel_data(self, start_year: str, start_month: str, end_year: str, end_month: str,
                       hour: str, minute: str) -> pd.DataFrame:
        """在一个指定的时间段内,获取每一天指定时刻的数据.

        由于每个月的天数不一样,因此,以下个月的1号为判定日期来确定时间段的终止年月,如有必要,再对数据进行删减处理.

        :param start_year: 时间段的起始年
        :param start_month: 时间段的起始月
        :param end_year: 时间段的终止年
        :param end_month: 时间段的终止月
        :param hour: 指定时刻的小时
        :param minute: 指定时刻的分钟
        :return:
        """
        start = pd.Timestamp(start_year + '-' + start_month + '-1')
        end = pd.Timestamp(end_year + '-' + end_month + '-1')
        df = self.data.truncate(before=start, after=end)
        return df[(df.hour == int(hour)) & (df.minute == int(minute))]


if __name__ == '__main__':
    warnings.filterwarnings('ignore')
    file_path = './data/opsd-time_series-2020-10-06/time_series_15min_singleindex.csv'
    ds = DataSet(file_path)
    h1, m1 = '1', '0'
    h2, m2 = '20', '0'
    df1 = ds.get_panel_data('2015', '1', '2016', '1', h1, m1)
    df2 = ds.get_panel_data('2015', '1', '2016', '1', h2, m2)
    x = df1.index
    y1 = df1['AT_price_day_ahead']
    y2 = df2['AT_price_day_ahead']

    fig, ax = plt.subplots()
    ax.set_title('价格曲线')
    ax.set_ylabel('价格')
    ax.set_xlabel('日期')
    ax.plot(x, y1, label=f'{h1}:{m1}价格')
    ax.plot(x, y2, label=f'{h2}:{m2}价格')
    ax.legend()
    plt.show()
